#!/usr/bin/env python

target_host = '10.64.70.25'
target_port = 8080

client_name = "client-1"

encrypt_url = '/sample/encrypt'
decrypt_url = '/sample/decrypt'
key_set_url = '/sample/key'

static_fetch_conf = {
	'dest_host' : client_name,
	'fetch_size': 1024 * 1024 * 4,
	'fetch_num' : 1024,
}

static_meta_conf = {
	'generator_name': client_name,
	'generator_cat' : 'frontend',
	'host_name'     : '192.168.1.1',
	'host_port'     : 8080,
}

#############################
meta_template_encrypt = '''<meta-data>
<task-id-chain>
<task-id>%%s</task-id>
</task-id-chain>
<journey>
<node>
<name>%(generator_name)s</name>
<in-time>%%u</in-time>
<out-time>%%u</out-time>
<service>%(generator_cat)s</service>
<host>%(host_name)s</host>
<port>%(host_port)u</port>
<logs/>
</node>
</journey>
<user-data/>
</meta-data>''' % static_meta_conf

meta_template_decrypt = '''<meta-data>
<task-id-chain>
<task-id>%%s</task-id>
</task-id-chain>
<journey>
<node>
<name>%(generator_name)s</name>
<in-time>%%u</in-time>
<out-time>%%u</out-time>
<service>%(generator_cat)s</service>
<host>%(host_name)s</host>
<port>%(host_port)u</port>
<logs/>
</node>
</journey>
<user-data>
<info><encrypted>1</encrypted></info>
<configure>
<key>%%s</key>
<method>default</method>
</configure>
</user-data>
</meta-data>''' % static_meta_conf

fetch_template = '''
<p1:common_pop_request xmlns:p1="http://www.trendmicro.com/nebula/xml_schema">
<pop_count>%(fetch_num)u</pop_count>
<pop_size>%(fetch_size)u</pop_size>
<dest_host>%(dest_host)s</dest_host>
</p1:common_pop_request>
''' % static_fetch_conf

data_flow_data_template = '%s'

import httplib
import itertools
import math
import os
import sys
import time
import StringIO

try:
	from uuid import uuid4 as _uuid4
	def uuid():
		return '%s' % _uuid4()
except ImportError:
	import random
	def _rand_4c():
		return '%04x' % random.randint(0, 0xFFFF)
	def uuid():
		return '%s%s-%s-%s-%s-%s%s%s' % tuple(_rand_4c() for x in range(8))

def push_encrypt_tasks(tasks, task_data, feed_url):
	dummy_id = uuid()
	body_length = (8 + len(data_flow_data_template % dummy_id) + 8 + 4 + 8 +
			len(meta_template_encrypt % (dummy_id, int(time.time() - 1.0), int(time.time()))) 
			+ 4 + 8 + len(task_data)) * tasks + 8
	del dummy_id

	conn = httplib.HTTPConnection(target_host, target_port)
	conn.connect()
	conn.putrequest('POST', feed_url)
	conn.putheader('Content-Length', body_length)
	conn.endheaders()

	for i in xrange(tasks):
		task_id = uuid()
		data_flow_data = data_flow_data_template % task_id
		#assuming that the task is generated 1.0 second ago and leave former node now
		meta = meta_template_encrypt % (task_id, int(time.time() - 1.0), int(time.time()))
		data = task_data

		conn.send('%08X' % len(data_flow_data))
		conn.send(data_flow_data)
		conn.send('%08X' % 2)
		conn.send('META')
		conn.send('%08X' % len(meta))
		conn.send(meta)
		conn.send('TASK')
		conn.send('%08X' % len(data))
		conn.send(data)
	conn.send('FFFFFFFF')
	resp  = conn.getresponse()
	reply = resp.read()
	return '%d %s %s' % (resp.status, resp.msg, reply)

def push_decrypt_tasks(tasks, task_data, key, feed_url):
	dummy_id = uuid()
	body_length = (8 + len(data_flow_data_template % dummy_id) + 8 + 4 + 8 +
			len(meta_template_decrypt % (dummy_id, int(time.time() - 1.0), int(time.time()), key)) 
			+ 4 + 8 + len(task_data)) * tasks + 8
	del dummy_id

	conn = httplib.HTTPConnection(target_host, target_port)
	conn.connect()
	conn.putrequest('POST', feed_url)
	conn.putheader('Content-Length', body_length)
	conn.endheaders()

	for i in xrange(tasks):
		task_id = uuid()
		data_flow_data = data_flow_data_template % task_id
		#assuming that the task is generated 1.0 second ago and leave former node now
		meta = meta_template_decrypt % (task_id, int(time.time() - 1.0), int(time.time()), key)
		data = task_data

		conn.send('%08X' % len(data_flow_data))
		conn.send(data_flow_data)
		conn.send('%08X' % 2)
		conn.send('META')
		conn.send('%08X' % len(meta))
		conn.send(meta)
		conn.send('TASK')
		conn.send('%08X' % len(data))
		conn.send(data)
	conn.send('FFFFFFFF')
	resp  = conn.getresponse()
	reply = resp.read()
	return '%d %s %s' % (resp.status, resp.msg, reply)

def fetch_tasks(fetch_url):
	data = fetch_template
	conn = httplib.HTTPConnection(target_host, target_port)
	conn.connect()
	conn.putrequest('PUT', fetch_url)
	conn.putheader('Content-Length', len(data))
	conn.endheaders()
	conn.send(data)
	resp = conn.getresponse()	
	parse_task_stream(resp)
	return (resp.status, resp.msg)

def parse_task_stream(stream):

    try:
        while True:
            task_id_len = int(stream.read(8), 16)
            if task_id_len == 0xFFFFFFFF:
                break

            task_id = stream.read(task_id_len)

            assert stream.read(8) == '00000002'
            assert stream.read(4) == 'META'

            meta_len = int(stream.read(8), 16)
            meta     = stream.read(meta_len)

            assert stream.read(4) == 'TASK'

            data_len = int(stream.read(8), 16)
            data     = stream.read(data_len)

            open(task_id + '.meta', 'wb').write(meta)
            open(task_id + '.task', 'wb').write(data)

            print 'task has been saved to %s + %s' % (task_id + '.meta', task_id + '.task')

    except ValueError:
        print 'Task data corrupted!'
        pass

def update_key(key):
	data = key
	conn = httplib.HTTPConnection(target_host, target_port)
	conn.connect()
	conn.putrequest('POST', key_set_url)
	conn.putheader('Content-Length', len(data))
	conn.endheaders()
	conn.send(data)
	resp = conn.getresponse()	
	resp.read()
	return '%d %s' % (resp.status, resp.msg)


def print_usage():
    print '''
usage: sample-client.py <operation [parameters] >
    available operations:
    	-e  <file name>        : encrypt a file
        -d  <file name> <key>  : decrypt a file, key must be offerred
        -er                    : get encrypt result, save to files
        -dr                    : get decrypt result, save to files
        -k  <key>              : update encrypt key specified        
        '''

def main():
	args = sys.argv[1:]

	if len(args) == 0:
		print_usage()
		return

	if args[0] == '-e' and len(args) == 2:
		if not os.path.exists(args[1]):
			print 'file %s not exists!'
			return

		data = open(args[1]).read()
		push_encrypt_tasks(1, data, encrypt_url)
	
	elif (args[0] == '-d') and len(args) == 3:
		if not os.path.exists(args[1]):
			print 'file %s not exists!'
			return
		data = open(args[1]).read()
		push_decrypt_tasks(1, data, args[2], decrypt_url)

	elif (args[0] == '-er'):
		fetch_tasks(encrypt_url)
	elif (args[0] == '-dr'):
		fetch_tasks(decrypt_url)
	elif (args[0] == '-k' and args[1]):
		update_key(args[1])
	else:
		print_usage()
		return
	

if __name__ == '__main__':
	try:
		main()
	except KeyboardInterrupt:
		pass

